# -*- coding: utf-8 -*-
import sys
import subprocess
import os
import argparse
import glob

def run_script(script_name, stock_code, frequency):
    """
    运行一个指定的子脚本, 并传递股票代码和时间周期作为参数.
    """
    script_path = script_name
    if not os.path.exists(script_path):
        print(f"错误：找不到脚本文件 '{script_path}'.")
        sys.exit(1)
        
    command = [sys.executable, script_path, stock_code, frequency]
    
    print(f"\n{'='*20}")
    print(f"  执行脚本: {script_name} (股票: {stock_code}, 周期: {frequency})")
    print(f"{'='*20}")
    
    process = subprocess.Popen(command, stdout=subprocess.PIPE, stderr=subprocess.STDOUT, text=True, encoding='utf-8')
    
    while True:
        output = process.stdout.readline()
        if output == '' and process.poll() is not None:
            break
        if output:
            print(output.strip())
            
    if process.returncode != 0:
        print(f"\n错误：脚本 '{script_name}' 执行失败. 流程终止.")
        sys.exit(1)

def clean_generated_files():
    """
    清理所有由本系统生成的 csv, png, json/h5, 和 pkl 文件.
    """
    print("开始清理生成的产物文件...")
    
    files_to_delete = []
    # 清理output目录中的所有文件
    output_patterns = [
        'output/data/stock_data_*.csv',
        'output/plots/prediction_vs_actual_*.png',
        'output/plots/transformer_prediction_*.png',  # Transformer预测图表
        'output/plots/tcn_training_history_*.png',    # TCN训练历史图表
        'output/models/xgb_model_*.json',
        'output/models/lstm_model_*.h5',
        'output/models/rf_model_*.pkl',              # RandomForest模型文件
        'output/models/tcn_model_*.h5',              # TCN模型文件
        'output/models/transformer_*.weights.h5',    # Transformer模型文件
        'output/scalers/tokenizer_*.pkl',            # Tokenizer文件
        'output/scalers/price_scaler_*.pkl',         # 价格归一化器
        'output/scalers/tcn_*_scaler_*.pkl',         # TCN归一化器文件
        'output/scalers/*_scaler_*.pkl'              # 其他归一化器文件 (X_scaler_*, y_scaler_*)
    ]
    
    # 也清理根目录的旧文件（兼容性）
    root_patterns = [
        'stock_data_*.csv',
        'prediction_vs_actual_*.png',
        'transformer_prediction_*.png',
        'tcn_training_history_*.png',
        'xgb_model_*.json',
        'lstm_model_*.h5',
        'rf_model_*.pkl',
        'tcn_model_*.h5',
        'transformer_*.weights.h5',
        'tokenizer_*.pkl',
        'price_scaler_*.pkl',
        'tcn_*_scaler_*.pkl',
        '*_scaler_*.pkl'
    ]
    
    all_patterns = output_patterns + root_patterns
    
    for pattern in all_patterns:
        files_to_delete.extend(glob.glob(pattern))
        
    if not files_to_delete:
        print("没有找到任何需要清理的文件.")
        return
        
    print("将要删除以下文件:")
    for f in files_to_delete:
        print(f" - {f}")
        
    confirm = input("你确定要删除这些文件吗？ (y/n): ")
    if confirm.lower() == 'y':
        for f in files_to_delete:
            try:
                os.remove(f)
                print(f"已删除: {f}")
            except OSError as e:
                print(f"删除文件时出错: {f}, 错误: {e}")
        print("\n清理完成！")
    else:
        print("操作已取消")

def add_common_arguments(parser):
    "一个辅助函数, 为给定的解析器添加通用参数."
    parser.add_argument('stock_code', type=str, help="股票代码")
    parser.add_argument('-f', '--frequency', choices=['daily', 'hourly'], default='daily', help="时间周期")

def main():
    # --- 1. 使用 argparse 解析命令行参数 ---
    parser = argparse.ArgumentParser(description="A-Share Stock Technical Analysis and Prediction Tool")
    subparsers = parser.add_subparsers(dest='action', help='Available actions')
    subparsers.required = True

    # 'train' 命令 (XGBoost)
    parser_train = subparsers.add_parser('train', help="训练 XGBoost 模型并立即预测")
    add_common_arguments(parser_train)

    # 'train_lstm' 命令
    parser_lstm = subparsers.add_parser('train_lstm', help="训练 LSTM 模型并立即预测")
    add_common_arguments(parser_lstm)

    # 'train_rf' 命令 (RandomForest)
    parser_rf = subparsers.add_parser('train_rf', help="训练 RandomForest 模型并立即预测")
    add_common_arguments(parser_rf)

    # 'train_tcn' 命令 (TCN - Temporal Convolutional Network)
    parser_tcn = subparsers.add_parser('train_tcn', help="训练 TCN (Google验证) 模型并立即预测")
    add_common_arguments(parser_tcn)

    # 'predict' 命令 (XGBoost)
    parser_predict = subparsers.add_parser('predict', help="使用已有的 XGBoost 模型进行预测")
    add_common_arguments(parser_predict)

    # 'predict_lstm' 命令
    parser_predict_lstm = subparsers.add_parser('predict_lstm', help="使用已有的 LSTM 模型进行预测")
    add_common_arguments(parser_predict_lstm)

    # 'predict_rf' 命令 (RandomForest)
    parser_predict_rf = subparsers.add_parser('predict_rf', help="使用已有的 RandomForest 模型进行预测")
    add_common_arguments(parser_predict_rf)

    # 'predict_tcn' 命令 (TCN)
    parser_predict_tcn = subparsers.add_parser('predict_tcn', help="使用已有的 TCN 模型进行预测")
    add_common_arguments(parser_predict_tcn)

    # 'train_transformer' 命令 (Stock-GPT Phase 1)
    parser_transformer = subparsers.add_parser('train_transformer', help="训练 Stock-GPT Transformer 模型 (Phase 1)")
    add_common_arguments(parser_transformer)
    parser_transformer.add_argument('--force-cpu', action='store_true', help='强制使用CPU模式')
    parser_transformer.add_argument('--epochs', type=int, default=50, help='训练轮数')

    # 'predict_transformer' 命令
    parser_predict_transformer = subparsers.add_parser('predict_transformer', help="使用已有的 Transformer 模型进行预测")
    add_common_arguments(parser_predict_transformer)
    parser_predict_transformer.add_argument('--force-cpu', action='store_true', help='强制使用CPU模式')

    # 'clean' 命令
    parser_clean = subparsers.add_parser('clean', help="清理所有生成的产物文件")

    args = parser.parse_args()
    
    # --- 2. 根据动作执行相应流程 ---
    if args.action in ['train', 'train_lstm', 'train_rf', 'train_tcn', 'train_transformer', 'predict', 'predict_lstm', 'predict_rf', 'predict_tcn', 'predict_transformer']:
        stock_code = args.stock_code
        frequency = args.frequency
        
        if args.action == 'train':
            print(f"=== 开始为股票 {stock_code} 执行完整的 {frequency} 周期 XGBoost 训练流程 ===")
            run_script('01_fetch_data.py', stock_code, frequency)
            run_script('02_preprocess_data.py', stock_code, frequency)
            run_script('03_calculate_indicators.py', stock_code, frequency)
            run_script('04_train_model.py', stock_code, frequency)
            print(f"\n--- 训练完成，自动执行一次初始预测 ---")
            run_script('05_predict_next.py', stock_code, frequency)
            print(f"\n🎉🎉🎉 XGBoost 训练和初始预测流程全部完成！ 🎉🎉🎉")

        elif args.action == 'train_lstm':
            print(f"=== 开始为股票 {stock_code} 执行 {frequency} 周期 LSTM 训练流程 ===")
            run_script('01_fetch_data.py', stock_code, frequency)
            run_script('02_preprocess_data.py', stock_code, frequency)
            run_script('03_calculate_indicators.py', stock_code, frequency)
            run_script('04b_train_lstm.py', stock_code, frequency)
            print(f"\n--- 训练完成，自动执行一次初始预测 ---")
            run_script('05b_predict_lstm.py', stock_code, frequency)
            print(f"\n🎉🎉🎉 LSTM 训练和初始预测流程全部完成！ 🎉🎉🎉")

        elif args.action == 'train_rf':
            print(f"=== 开始为股票 {stock_code} 执行 {frequency} 周期 RandomForest 训练流程 ===")
            run_script('01_fetch_data.py', stock_code, frequency)
            run_script('02_preprocess_data.py', stock_code, frequency)
            run_script('03_calculate_indicators.py', stock_code, frequency)
            run_script('04c_train_randomforest.py', stock_code, frequency)
            print(f"\n--- 训练完成，自动执行一次初始预测 ---")
            run_script('05c_predict_randomforest.py', stock_code, frequency)
            print(f"\n🎉🎉🎉 RandomForest 训练和初始预测流程全部完成！ 🎉🎉🎉")

        elif args.action == 'train_tcn':
            print(f"=== 开始为股票 {stock_code} 执行 {frequency} 周期 TCN (Google验证) 训练流程 ===")
            run_script('01_fetch_data.py', stock_code, frequency)
            run_script('02_preprocess_data.py', stock_code, frequency)
            run_script('03_calculate_indicators.py', stock_code, frequency)
            run_script('04d_train_tcn.py', stock_code, frequency)
            print(f"\n--- 训练完成，自动执行一次初始预测 ---")
            run_script('05d_predict_tcn.py', stock_code, frequency)
            print(f"\n🎉🎉🎉 TCN 训练和初始预测流程全部完成！ 🎉🎉🎉")

        elif args.action == 'predict':
            run_script('05_predict_next.py', stock_code, frequency)
            print(f"\n🚀 XGBoost 预测流程完成！ 🚀")
            
        elif args.action == 'predict_lstm':
            run_script('05b_predict_lstm.py', stock_code, frequency)
            print(f"\n🚀 LSTM 预测流程完成！ 🚀")

        elif args.action == 'predict_rf':
            run_script('05c_predict_randomforest.py', stock_code, frequency)
            print(f"\n🚀 RandomForest 预测流程完成！ 🚀")

        elif args.action == 'predict_tcn':
            run_script('05d_predict_tcn.py', stock_code, frequency)
            print(f"\n🚀 TCN 预测流程完成！ 🚀")
            
        elif args.action == 'train_transformer':
            print(f"=== 开始为股票 {stock_code} 执行 {frequency} 周期 Stock-GPT Transformer 训练流程 ===")
            run_script('01_fetch_data.py', stock_code, frequency)
            run_script('02_preprocess_data.py', stock_code, frequency)
            run_script('03_calculate_indicators.py', stock_code, frequency)
            
            # 构建Transformer训练命令
            transformer_cmd = [sys.executable, '06_train_transformer.py', stock_code, '-f', frequency]
            if hasattr(args, 'force_cpu') and args.force_cpu:
                transformer_cmd.append('--force-cpu')
            if hasattr(args, 'epochs') and args.epochs:
                transformer_cmd.extend(['--epochs', str(args.epochs)])
                
            print(f"\n{'='*20}")
            print(f"  执行脚本: 06_train_transformer.py (股票: {stock_code}, 周期: {frequency})")
            print(f"{'='*20}")
            
            process = subprocess.Popen(transformer_cmd, stdout=subprocess.PIPE, stderr=subprocess.STDOUT, text=True, encoding='utf-8')
            
            while True:
                output = process.stdout.readline()
                if output == '' and process.poll() is not None:
                    break
                if output:
                    print(output.strip())
                    
            if process.returncode != 0:
                print(f"\n错误：Transformer训练失败. 流程终止.")
                sys.exit(1)
                
            print(f"\n--- 训练完成，自动执行一次初始预测 ---")
            
            # 构建预测命令
            predict_cmd = [sys.executable, '06b_predict_transformer.py', stock_code, '-f', frequency]
            if hasattr(args, 'force_cpu') and args.force_cpu:
                predict_cmd.append('--force-cpu')
                
            print(f"\n{'='*20}")
            print(f"  执行脚本: 06b_predict_transformer.py (股票: {stock_code}, 周期: {frequency})")
            print(f"{'='*20}")
            
            process = subprocess.Popen(predict_cmd, stdout=subprocess.PIPE, stderr=subprocess.STDOUT, text=True, encoding='utf-8')
            
            while True:
                output = process.stdout.readline()
                if output == '' and process.poll() is not None:
                    break
                if output:
                    print(output.strip())
                    
            print(f"\n🎉🎉🎉 Stock-GPT Transformer 训练和初始预测流程全部完成！ 🎉🎉🎉")
            
        elif args.action == 'predict_transformer':
            # 构建预测命令
            predict_cmd = [sys.executable, '06b_predict_transformer.py', stock_code, '-f', frequency]
            if hasattr(args, 'force_cpu') and args.force_cpu:
                predict_cmd.append('--force-cpu')
                
            print(f"\n{'='*20}")
            print(f"  执行脚本: 06b_predict_transformer.py (股票: {stock_code}, 周期: {frequency})")
            print(f"{'='*20}")
            
            process = subprocess.Popen(predict_cmd, stdout=subprocess.PIPE, stderr=subprocess.STDOUT, text=True, encoding='utf-8')
            
            while True:
                output = process.stdout.readline()
                if output == '' and process.poll() is not None:
                    break
                if output:
                    print(output.strip())
                    
            print(f"\n🚀 Stock-GPT Transformer 预测流程完成！ 🚀")
            
    elif args.action == 'clean':
        clean_generated_files()

if __name__ == "__main__":
    main()
